from openai import OpenAI
import time

from group_profile.config_util import ConfigUtil

"""封装了大模型接口的基本问答智能体
"""





class AIAgent:


    # deepseek r1
    client = OpenAI(

        # deepseek apis
        base_url="https://api.deepseek.com",
        api_key=ConfigUtil.get_api_key(),

        # openrouter apis
        # base_url="https://openrouter.ai/api/v1",
        # api_key="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",
        # api_key="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",
    )

    def complete(self, prompt):
        # 记录请求发送前的时间
        start_time = time.time()
        response = self.client.chat.completions.create(
            extra_headers={
                "HTTP-Referer": "<YOUR_SITE_URL>",  # Optional. Site URL for rankings on openrouter.ai.
                "X-Title": "<YOUR_SITE_NAME>",  # Optional. Site title for rankings on openrouter.ai.
            },
            # model="deepseek/deepseek-r1:free",
            model="deepseek-chat",
            messages=[
                {
                    "role": "system",
                    "content": self.get_system_prompt()
                },

                {
                    "role": "user",
                    "content": prompt
                }
            ],
            stream=False
        )
        return response

    def stream_complete(self, prompt, history, system_prompt=None):
        if not system_prompt:
            system_prompt = self.get_system_prompt()
        messages = [
            {
                "role": "system",
                "content": system_prompt
            }
        ]
        if history:
            for h in history[-5:]:  # 仅考虑最近5次对话
                messages.append(h)
        messages.append({
            "role": "user",
            "content": prompt
        })
        # logging.info(f"发送message：{messages}")
        response = self.client.chat.completions.create(
            extra_headers={
                "HTTP-Referer": "<YOUR_SITE_URL>",  # Optional. Site URL for rankings on openrouter.ai.
                "X-Title": "<YOUR_SITE_NAME>",  # Optional. Site title for rankings on openrouter.ai.
            },
            model="deepseek-chat",
            messages=messages,
            stream=True
        )
        return response

    def get_system_prompt(self):
        return "你是洞察基础AI，在每一个回答的前面都先介绍自己"


if __name__ == '__main__':
    agent = AIAgent()
    # 记录请求发送前的时间
    start_time = time.time()
    response1 = agent.stream_complete("写一首小诗，不超过10个字", [])
    # 创建变量来收集事件流
    collected_chunks = []
    collected_messages = []
    # 通过事件流迭代
    for chunk in response1:
        chunk_time = time.time() - start_time
        collected_chunks.append(chunk)
        chunk_message = chunk.choices[0].delta
        collected_messages.append(chunk_message)
        # print(f"Message received {chunk_time:.2f} seconds after request: {chunk_message}")
        print(chunk.choices[0].delta.content, end='')

    print()

    # 打印完整响应和文本
    # full_reply_content = ''.join([m.content for m in collected_messages])
    # print(f"Full conversation received: {full_reply_content}")
